#include "udp_server.h"
#include <arpa/inet.h>
#include <iostream>
#include "buffer.h"
#include "data_segment.h"
#include "packet_statistics.h"
#include "sliding_window.h"
#include "spdlog/spdlog.h"

using namespace std;
using namespace safe_udp;

UdpServer::UdpServer() {
    sliding_window_ = make_unique<SlidingWindow>();
    packet_statistics_ = make_unique<PacketStatistics>();

    sockfd_ = 0;
    smoothed_rtt_ = 20000;
    smoothed_timeout_ = 30000;
    dev_rtt_ = 0;

    initial_seq_number_ = 67;
    start_byte_ = 0;

    ssthresh_ = 128;
    cwnd_ = 1;
    is_slow_start_ = true;
    is_cong_avd_ = false;
    is_fast_recovery_ = false;
}

UdpServer::~UdpServer() {
    close(sockfd_);
    file_.close();
}

char* UdpServer::GetRequest(int client_sockfd) {
    char* buffer = reinterpret_cast<char*>(calloc(MAX_PACKET_SIZE, sizeof(char)));
    struct sockaddr_in client_address;
    socklen_t addr_size;
    memset(buffer, 0, MAX_PACKET_SIZE);
    addr_size = sizeof(client_address);
    recvfrom(client_sockfd, buffer, MAX_PACKET_SIZE, 0, (struct sockaddr*)&client_address, &addr_size);
    cli_address_ = client_address;
    return buffer;
}

bool UdpServer::Openfile(const std::string& file_name) {
    spdlog::info("Opening the file {}", file_name);
    file_.open(file_name, std::ios::in);
    if (!file_.is_open()) {
        spdlog::warn("File {} opening failed", file_name);
        return false;
    } else {
        spdlog::info("File {} opening success", file_name);
        return true;
    }
}

void UdpServer::StartFileTransfer() {
    spdlog::info("Starting the file transfer");
    file_.seekg(0, std::ios::end);
    file_length_ = file_.tellg();
    file_.seekg(0, std::ios::beg);

    send();
}

void UdpServer::SendError() {
    string error("FILE NOT FOUND");
    sendto(sockfd_, error.c_str(), error.size(), 0, (struct sockaddr*)&cli_address_, sizeof(cli_address_));
}

int UdpServer::StartServer(int port) {
    int sfd;
    struct sockaddr_in server_addr;
    spdlog::info("Starting the webserver... port: {}", port);
    sfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sfd < 0) {
        spdlog::error("Faild to socket !!!");
        exit(0);
    }
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr("127.0.0.1");
    server_addr.sin_port = htons(port);

    if (bind(sfd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        spdlog::error("binding error !!!");
        exit(0);
    }
    spdlog::info("Server Bind set to addr: {}", inet_ntoa(server_addr.sin_addr));
    spdlog::info("Server Bind set to port: {}", port);
    spdlog::info("Server Bind set to family: {}", server_addr.sin_family);
    spdlog::info("Started successfully");
    sockfd_ = sfd;
    return sfd;
}

void UdpServer::send() {
    spdlog::info("Entering Send()");
    int sent_count = 1;
    int sent_count_limit = 1;

    struct timeval process_start_time;
    struct timeval process_end_time;
    gettimeofday(&process_start_time, NULL);

    if (sliding_window_->last_packet_sent_ == -1) {
        start_byte_ = 0;
    }
    while (start_byte_ <= file_length_) {
        fd_set rfds;
        struct timeval tv;
        int res;

        sent_count = 1;
        sent_count_limit = std::min(rwnd_, cwnd_);

        spdlog::info("SEND START !!!!");
        spdlog::info("Before the window rwnd_: {} cwnd_: {} window used: {}", rwnd_, cwnd_, sliding_window_->last_packet_sent_ - sliding_window_->last_acked_packet_);
        while (sliding_window_->last_packet_sent_ - sliding_window_->last_acked_packet_ <= std::min(rwnd_, cwnd_) &&
               sent_count <= sent_count_limit) {
            send_packed(start_byte_ + initial_seq_number_, start_byte_);
            if (is_slow_start_) {
                packet_statistics_->slow_start_packet_sent_count_++;
            } else if (is_cong_avd_) {
                packet_statistics_->cong_avd_packet_sent_count_++;
            }

            start_byte_ = start_byte_ + MAX_DATA_SIZE;
            if (start_byte_ > file_length_) {
                spdlog::info("No more data left to be sent");
                break;
            }
            sent_count++;
        }
        spdlog::info("send end !!!");

        // socket listen whith  timeout
        FD_ZERO(&rfds);
        FD_SET(sockfd_, &rfds);

        tv.tv_sec = (int64_t)smoothed_timeout_ / 1000000;
        tv.tv_usec = smoothed_timeout_ - tv.tv_sec;

        spdlog::info(" SELECT START: {} !!!", smoothed_timeout_);
        while (true) {
            res = select(sockfd_ + 1, &rfds, NULL, NULL, &tv);
            if (res == -1) {
                spdlog::error("Error in select");
            } else if (res > 0) {
                wait_for_ack();
                if (cwnd_ >= ssthresh_) {
                    // 慢启动 ---> 拥塞避免
                    spdlog::info("CHANGE TO CONG AVD");
                    is_cong_avd_ = true;
                    is_slow_start_ = false;

                    cwnd_ = 1;
                    ssthresh_ = 64;
                }

                if (sliding_window_->last_acked_packet_ == sliding_window_->last_packet_sent_) {
                    if (is_slow_start_) {
                        cwnd_ += 2;
                    } else {
                        cwnd_ += 1;
                    }
                    break;
                }
            } else {
                // 拥塞发生 --- 超时重传
                spdlog::info("Timeout occurred SELECT:: {}", smoothed_timeout_);
                ssthresh_ = cwnd_ / 2;
                if (ssthresh_ < 1) {
                    ssthresh_ = 1;
                }
                cwnd_ = 1;

                // 重新开始慢启动
                if (is_fast_recovery_) {
                    is_fast_recovery_ = false;
                }
                is_slow_start_ = true;
                is_cong_avd_ = false;

                // retransmit all unacked segments
                for (int i = sliding_window_->last_acked_packet_ + 1;
                     i < sliding_window_->last_packet_sent_; ++i) {
                    int retransmit_start_byte = 0;
                    if (sliding_window_->last_acked_packet_ != -1) {
                        retransmit_start_byte =
                            sliding_window_
                                ->sliding_window_buffers_[sliding_window_
                                                              ->last_acked_packet_]
                                .first_byte_ +
                            MAX_DATA_SIZE;
                    }
                    spdlog::info("Timeout Retransmit seq number: {}", retransmit_start_byte + initial_seq_number_);
                    retransmit_segment(retransmit_start_byte);
                    packet_statistics_->retransmit_count_++;
                    spdlog::info("Timeout: retransmission at {}", retransmit_start_byte);
                }
                break;
            }
        }
        spdlog::info("SELECT END !!!");
        spdlog::info("current byte:: {} length {}", start_byte_, file_length_);
    }
    gettimeofday(&process_end_time, NULL);
    int64_t total_time = (process_end_time.tv_sec * 1000000 + process_end_time.tv_usec) -
                         (process_start_time.tv_sec * 1000000 + process_start_time.tv_usec);
    int total_packet_sent = packet_statistics_->slow_start_packet_sent_count_ +
                            packet_statistics_->cong_avd_packet_sent_count_;

    spdlog::info("\n========================================");
    spdlog::info("Total time: {} secs", (float)total_time / 1000000);
    spdlog::info("Statistics: 拥塞控制--慢启动: {}", packet_statistics_->slow_start_packet_sent_count_);
    spdlog::info("拥塞控制--拥塞避免: {}", packet_statistics_->cong_avd_packet_sent_count_);
    spdlog::info("Statistics: Slow start: {}", ((float)packet_statistics_->slow_start_packet_sent_count_ / total_packet_sent) * 100);
    spdlog::info("% CongAvd {} %", ((float)packet_statistics_->cong_avd_packet_sent_count_ / total_packet_sent) * 100);
    spdlog::info("Statistics: Retransmissions: {}", packet_statistics_->retransmit_count_);
    spdlog::info("========================================");
}

void UdpServer::send_packed(int seq_number, int start_byte) {
    bool last_packet = false;
    int data_length = 0;
    if (file_length_ <= start_byte + MAX_DATA_SIZE) {
        spdlog::info("Last packet to be sent !!!");
        data_length = file_length_ - start_byte;
        last_packet = true;
    } else {
        data_length = MAX_DATA_SIZE;
    }

    struct timeval time;
    gettimeofday(&time, NULL);
    if (sliding_window_->last_packet_sent_ != -1 &&
        start_byte < sliding_window_->sliding_window_buffers_[sliding_window_->last_packet_sent_].first_byte_) {
        for (int i = sliding_window_->last_acked_packet_ + 1; i < sliding_window_->last_packet_sent_; ++i) {
            if (sliding_window_->sliding_window_buffers_[i].first_byte_ == start_byte) {
                sliding_window_->sliding_window_buffers_[i].time_sent_ = time;
                break;
            }
        }
    } else {
        SlidWinBuffer slidingWindowBuffer;
        slidingWindowBuffer.first_byte_ = start_byte;
        slidingWindowBuffer.data_length_ = data_length;
        slidingWindowBuffer.seq_num_ = initial_seq_number_ + start_byte;

        gettimeofday(&slidingWindowBuffer.time_sent_, NULL);
        sliding_window_->last_packet_sent_ = sliding_window_->AddToBuffer(slidingWindowBuffer);
    }
    read_file_and_send(last_packet, start_byte, start_byte + data_length);
}

void UdpServer::calculate_rtt_and_time(struct timeval start_time, struct timeval end_time) {
    if (start_time.tv_sec == 0 && start_time.tv_usec == 0) {
        return;
    }
    long sample_rtt = (end_time.tv_sec - start_time.tv_sec) * 1000000 + (end_time.tv_usec - start_time.tv_usec);
    smoothed_rtt_ = smoothed_rtt_ + 0.125 * (sample_rtt - smoothed_rtt_);
    dev_rtt_ = 0.75 * dev_rtt_ + 0.25 * (abs(smoothed_rtt_ - sample_rtt));
    smoothed_rtt_ = smoothed_rtt_ + 4 * dev_rtt_;

    if (smoothed_timeout_ > 1000000) {
        smoothed_timeout_ = rand() % 30000;
    }
}

void UdpServer::retransmit_segment(int index_number) {
    for (int i = sliding_window_->last_acked_packet_ + 1; i < sliding_window_->last_packet_sent_; ++i) {
        if (sliding_window_->sliding_window_buffers_[i].first_byte_ == index_number) {
            struct timeval time;
            gettimeofday(&time, NULL);
            sliding_window_->sliding_window_buffers_[i].time_sent_ = time;
            break;
        }
    }
    read_file_and_send(false, index_number, index_number + MAX_DATA_SIZE);
}

void UdpServer::read_file_and_send(bool fin_flag, int start_byte, int end_byte) {
    int data_length = end_byte - start_byte;
    if (file_length_ - start_byte < data_length) {
        data_length = file_length_ - start_byte;
        fin_flag = true;
    }
    char* fileData = reinterpret_cast<char*>(calloc(data_length, sizeof(char)));
    if (!file_.is_open()) {
        spdlog::error("File open failed !!!");
        return;
    }
    file_.seekg(start_byte);
    file_.read(fileData, data_length);

    DataSegment* data_segment = new DataSegment();
    data_segment->seq_number_ = start_byte + initial_seq_number_;
    data_segment->ack_number_ = 0;
    data_segment->ack_flag_ = false;
    data_segment->fin_flag_ = fin_flag;
    data_segment->data_ = fileData;

    send_data_segment(data_segment);
    spdlog::info("Packet sent seq number: {}", data_segment->seq_number_);
    free(fileData);
    free(data_segment);
}

void UdpServer::send_data_segment(DataSegment* data_segment) {
    char* datagramChars = data_segment->SerializeToCharArray();
    sendto(sockfd_, datagramChars, MAX_PACKET_SIZE, 0, (struct sockaddr*)&cli_address_, sizeof(cli_address_));
    free(datagramChars);
}

void UdpServer::wait_for_ack() {
    unsigned char buffer[MAX_PACKET_SIZE];
    memset(buffer, 0, MAX_PACKET_SIZE);
    socklen_t addr_size;
    struct sockaddr_in client_address;
    addr_size = sizeof(client_address);
    int n = 0;
    int ack_number;
    while (true) {
        n = recvfrom(sockfd_, buffer, MAX_PACKET_SIZE, 0, (struct sockaddr*)&client_address, &addr_size);
        if (n > 0) {
            break;
        }
    }
    DataSegment ack_segment;
    ack_segment.DeserializeToDataSegment(buffer, n);
    SlidWinBuffer last_packet_acked_buffer = sliding_window_->sliding_window_buffers_[sliding_window_->last_acked_packet_];
    if (ack_segment.ack_flag_) {
        if (ack_segment.ack_number_ == sliding_window_->send_base_) {
            spdlog::info("DUP ACK Received: ack_number: {}", ack_segment.ack_number_);
            sliding_window_->dup_ack_++;
            // 快速重传
            if (sliding_window_->dup_ack_ == 3) {
                packet_statistics_->retransmit_count_++;
                spdlog::info("Fast Retransmit seq_number: ", ack_segment.ack_number_);
                retransmit_segment(ack_segment.ack_number_ - initial_seq_number_);
                sliding_window_->dup_ack_ = 0;
                if (cwnd_ > 1) {
                    cwnd_ /= 2;
                }
                ssthresh_ = cwnd_;
                is_fast_recovery_ = true;
            }
        } else if (ack_segment.ack_number_ > sliding_window_->send_base_) {
            if (is_fast_recovery_) {
                cwnd_++;
                is_fast_recovery_ = false;
                is_cong_avd_ = true;
                is_slow_start_ = false;
            }

            sliding_window_->dup_ack_ = 0;
            sliding_window_->send_base_ = ack_segment.ack_number_;
            if (sliding_window_->last_acked_packet_ == -1) {
                sliding_window_->last_acked_packet_ = 0;
                last_packet_acked_buffer = sliding_window_->sliding_window_buffers_[sliding_window_->last_acked_packet_];
            }
            ack_number = last_packet_acked_buffer.seq_num_ + last_packet_acked_buffer.data_length_;
            while (ack_number < ack_segment.ack_number_) {
                sliding_window_->last_acked_packet_++;
                last_packet_acked_buffer = sliding_window_->sliding_window_buffers_[sliding_window_->last_acked_packet_];
                ack_number = last_packet_acked_buffer.seq_num_ + last_packet_acked_buffer.data_length_;
            }

            struct timeval startTime = last_packet_acked_buffer.time_sent_;
            struct timeval endTime;
            gettimeofday(&endTime, NULL);
            calculate_rtt_and_time(startTime, endTime);
        }
    }
}